package com.nbsaas.boot;

import com.baidu.aip.ocr.AipOcr;
import org.apache.poi.hwpf.HWPFDocument;
import org.apache.poi.hwpf.model.PicturesTable;
import org.apache.poi.hwpf.usermodel.Picture;
import org.apache.poi.xwpf.usermodel.XWPFDocument;
import org.apache.poi.xwpf.usermodel.XWPFPictureData;
import org.json.JSONArray;
import org.json.JSONObject;

import java.io.FileInputStream;
import java.io.IOException;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

/**
 * 提取 Word(DOC/DOCX) 内嵌图片并进行 OCR 识别，用于扫描件或图片型内容。
 */
public class WordImageOcrReader {

    private final AipOcr ocrClient;

    public WordImageOcrReader() {
        this.ocrClient = new AipOcr(OcrConfig.BAIDU_APP_ID, OcrConfig.BAIDU_API_KEY, OcrConfig.BAIDU_SECRET_KEY);
        this.ocrClient.setConnectionTimeoutInMillis(OcrConfig.OCR_CONNECTION_TIMEOUT);
        this.ocrClient.setSocketTimeoutInMillis(OcrConfig.OCR_SOCKET_TIMEOUT);
    }

    /**
     * 根据扩展名选择 DOCX 或 DOC 的图片提取与 OCR
     */
    public String extractText(Path wordPath) throws IOException {
        String name = wordPath.getFileName() != null ? wordPath.getFileName().toString().toLowerCase() : "";
        if (name.endsWith(".docx")) {
            return ocrDocxImages(wordPath);
        } else if (name.endsWith(".doc")) {
            return ocrDocImages(wordPath);
        } else {
            return "";
        }
    }

    private String ocrDocxImages(Path docxPath) throws IOException {
        List<String> pageTexts = new ArrayList<>();
        try (FileInputStream fis = new FileInputStream(docxPath.toFile());
             XWPFDocument doc = new XWPFDocument(fis)) {
            List<XWPFPictureData> pictures = doc.getAllPictures();
            System.out.println("启用 Word 图片 OCR 识别 (DOCX)，图片数量: " + pictures.size() + ", 文件: " + docxPath);
            for (XWPFPictureData pic : pictures) {
                byte[] data = pic.getData();
                String text = ocrImageBytes(data);
                if (!text.isBlank()) {
                    pageTexts.add(text);
                }
            }
        }
        return String.join("\n\n", pageTexts).trim();
    }

    private String ocrDocImages(Path docPath) throws IOException {
        List<String> pageTexts = new ArrayList<>();
        try (FileInputStream fis = new FileInputStream(docPath.toFile());
             HWPFDocument doc = new HWPFDocument(fis)) {
            PicturesTable picturesTable = doc.getPicturesTable();
            List<Picture> pictures = picturesTable.getAllPictures();
            System.out.println("启用 Word 图片 OCR 识别 (DOC)，图片数量: " + pictures.size() + ", 文件: " + docPath);
            for (Picture pic : pictures) {
                byte[] data = pic.getContent();
                String text = ocrImageBytes(data);
                if (!text.isBlank()) {
                    pageTexts.add(text);
                }
            }
        }
        return String.join("\n\n", pageTexts).trim();
    }

    private String ocrImageBytes(byte[] imageBytes) {
        HashMap<String, String> options = new HashMap<>();
        options.put("detect_direction", "true");
        options.put("probability", "true");
        JSONObject res = ocrClient.basicGeneral(imageBytes, options);
        if (res == null) {
            return "";
        }
        if (res.has("words_result")) {
            JSONArray arr = res.optJSONArray("words_result");
            if (arr != null) {
                StringBuilder sb = new StringBuilder();
                for (int i = 0; i < arr.length(); i++) {
                    JSONObject obj = arr.optJSONObject(i);
                    if (obj != null) {
                        String w = obj.optString("words", "");
                        if (!w.isEmpty()) {
                            sb.append(w).append("\n");
                        }
                    }
                }
                return sb.toString().trim();
            }
        }
        return res.optString("error_msg", "");
    }
}